#include "fiber.h"

static bool debug = false;

namespace zsylar {

// 当前线程上的协程控制信息
// 正在运行的协程
static thread_local Fiber* t_fiber = nullptr;
// 主协程
static thread_local std::shared_ptr<Fiber> t_thread_fiber = nullptr;
// 调度协程
static thread_local Fiber* t_scheduler_fiber = nullptr;

// 协程id
static std::atomic<uint64_t> s_fiber_id{0};
// 协程计数器
static std::atomic<uint64_t> s_fiber_count{0};

// 设置当前运行协程
void Fiber::SetThis(Fiber* f)
{
    t_fiber = f;
}

// 获取当前运行的协程
std::shared_ptr<Fiber> Fiber::GetThis()
{
    // 如果当前线程有运行的协程，返回该协程的智能指针
    if (t_fiber)
    {
        return t_fiber->shared_from_this();
    }

    // 如果当前线程没有运行的协程，创建主协程
    std::shared_ptr<Fiber> main_fiber(new Fiber());
    t_thread_fiber = main_fiber; // 设置主协程
    t_scheduler_fiber = main_fiber.get(); // 设置调度协程为主协程

    // 确保当前运行的协程是主协程
    assert(t_fiber == main_fiber.get());

    return t_fiber->shared_from_this(); // 返回主协程的智能指针
}

// 设置调度协程（默认为主协程）
void Fiber::SetSchedulerFiber(Fiber* f)
{
    t_scheduler_fiber = f;
}

// 得到当前运行的协程id
uint64_t Fiber::GetFiberId()
{
    // 如果当前线程有运行的协程，返回该协程的ID
    if (t_fiber)
    {
        return t_fiber->getId();
    }
    return (uint64_t)-1; // 如果没有运行的协程，返回无效ID
}

// 用于创建主协程（通常是程序的初始协程）
// 主协程是程序的初始协程，直接使用默认栈，不需要回调函数，状态为 RUNNING。
Fiber::Fiber()
{
    SetThis(this); // 设置当前协程为主协程
    m_state = RUNNING; // 设置状态为运行中

    // 获取当前上下文
    if (getcontext(&m_ctx))
    {
        std::cerr << "getcontext() failed\n";
        pthread_exit(NULL);
    }

    m_id = ++s_fiber_id; // 分配一个新的协程ID
    s_fiber_count++; // 增加协程计数器

    if (debug)
    {
        std::cout << "Fiber::Fiber(): main id=" << m_id << std::endl;
    }
}

// 用于创建子协程（通常是程序的初始协程）
// 子协程是用户定义的协程，需要显式分配栈空间，指定回调函数，状态为 READY，等待调度器运行。
Fiber::Fiber(std::function<void()> cb, size_t stacksize, bool run_in_scheduler)
    : m_cb(cb), m_runInScheduler(run_in_scheduler)
{
    m_state = READY; // 初始状态为准备就绪

    // 分配协程栈空间
    m_stacksize = stacksize ? stacksize : 128000; // 默认栈大小为128KB
    m_stack = malloc(m_stacksize); // 分配栈空间

    if (getcontext(&m_ctx))
    {
        std::cerr << "Fiber(std::function<void()> cb, size_t stacksize, bool run_in_scheduler) failed\n";
        pthread_exit(NULL);
    }

    m_ctx.uc_link = nullptr; // 设置上下文链接为nullptr
    m_ctx.uc_stack.ss_sp = m_stack; // 设置栈空间指针
    m_ctx.uc_stack.ss_size = m_stacksize; // 设置栈大小
    makecontext(&m_ctx, &Fiber::MainFunc, 0); // 设置协程函数入口

    m_id = ++s_fiber_id; // 分配一个新的协程ID
    s_fiber_count++; // 增加协程计数器

    if (debug)
    {
        std::cout << "Fiber::Fiber(): child id=" << m_id << std::endl;
    }
}

// 析构函数
Fiber::~Fiber()
{
    s_fiber_count--; // 减少协程计数器
    if (m_stack)
    {
        free(m_stack); // 释放栈空间
    }
    if (debug)
    {
        std::cout << "Fiber::~Fiber(): id=" << m_id << std::endl;
    }
}

// 重置一个协程
// 该函数用于重用一个已终止的协程，重新设置其状态、回调函数和栈空间。
// 在重置之前，确保协程的栈不为空且状态为已终止（TERM）。
// 重置后，协程的状态将变为准备就绪（READY），并且可以重新运行。
// 如果重置失败，将输出错误信息并终止线程。
// 注意：重置协程时，必须提供一个新的回调函数。
// 该函数会重新获取上下文，并设置协程的栈空间和入口函数。
// 如果获取上下文失败，将输出错误信息并终止线程。
// 重置协程后，可以通过调用 resume() 方法来恢复执行。
void Fiber::reset(std::function<void()> cb)
{
    assert(m_stack != nullptr && m_state == TERM); // 确保栈不为空且状态为已终止

    m_state = READY; // 重置状态为准备就绪
    m_cb = cb; // 设置新的协程函数

    if (getcontext(&m_ctx))
    {
        std::cerr << "Fiber::reset() failed\n";
        pthread_exit(NULL);
    }

    m_ctx.uc_link = nullptr; // 设置上下文链接为nullptr
    m_ctx.uc_stack.ss_sp = m_stack; // 设置栈空间指针
    m_ctx.uc_stack.ss_size = m_stacksize; // 设置栈大小
    makecontext(&m_ctx, &Fiber::MainFunc, 0); // 设置协程函数入口
}

// 恢复执行协程
void Fiber::resume()
{
    assert(m_state == READY); // 确保协程状态为准备就绪

    m_state = RUNNING; // 设置状态为运行中

    if (m_runInScheduler)
    {
        // 运行在调度器中：
        //     调度器负责管理多个协程的切换。
        //     当前协程需要切换到调度器的上下文，以便调度器继续运行其他协程。
        SetThis(this); // 设置当前协程为正在运行的协程
        if (swapcontext(&(t_scheduler_fiber->m_ctx), &m_ctx))
        {
            std::cerr << "Fiber::resume() to t_scheduler_fiber failed\n";
            pthread_exit(NULL);
        }
    }
    else
    {
        // 运行在线程中：
        //     当前协程直接运行在线程中，没有调度器参与。
        //     协程需要切换到线程的上下文，以便线程继续执行其他任务。
        SetThis(this); // 设置当前协程为正在运行的协程
        if (swapcontext(&(t_thread_fiber->m_ctx), &m_ctx))
        {
            std::cerr << "Fiber::resume() to t_thread_fiber failed\n";
            pthread_exit(NULL);
        }
    }
}

// 协程暂停
void Fiber::yield()
{
    assert(m_state == RUNNING || m_state == TERM); // 确保协程状态为运行中或已终止

    if (m_state != TERM)
    {
        m_state = READY; // 如果不是已终止状态，设置状态为准备就绪
    }

    if (m_runInScheduler)
    {
        // 运行在调度器中：
        SetThis(t_scheduler_fiber); // 设置当前协程为调度器协程
        if (swapcontext(&m_ctx, &(t_scheduler_fiber->m_ctx)))
        {
            std::cerr << "Fiber::yield() to t_scheduler_fiber failed\n";
            pthread_exit(NULL);
        }
    }
    else
    {
        // 运行在线程中：
        SetThis(t_thread_fiber.get()); // 设置当前协程为线程协程
        if (swapcontext(&m_ctx, &(t_thread_fiber->m_ctx)))
        {
            std::cerr << "Fiber::yield() to t_thread_fiber failed\n";
            pthread_exit(NULL);
        }
    }
}

// 协程入口函数
void Fiber::MainFunc()
{
    std::shared_ptr<Fiber> curr = GetThis(); // 获取当前协程的智能指针
	assert(curr != nullptr); // 确保当前协程不为空

    curr->m_cb(); // 执行协程的回调函数
    curr->m_cb = nullptr; // 清空回调函数，避免重复执行
    curr->m_state = TERM; // 设置协程状态为已终止

    // 运行完毕 -> 出让执行权
    auto raw_ptr = curr.get(); // 获取当前协程的原始指针
    curr.reset(); // 重置智能指针，释放资源
    raw_ptr->yield(); // 让出执行权，切换到其他协程或线程
}

}